# encoding: UTF-8

import hashlib
import zlib
import json
from time import sleep
from threading import Thread

import websocket


# OKCOIN网站
OKCOIN_CNY = 'wss://real.okcoin.cn:10440/websocket/okcoinapi'
OKCOIN_USD = 'wss://real.okcoin.com:10440/websocket/okcoinapi'

# 账户货币代码
CURRENCY_CNY = 'cny'
CURRENCY_USD = 'usd'

# 电子货币代码
SYMBOL_BTC = 'btc'
SYMBOL_LTC = 'ltc'

# 行情深度
DEPTH_20 = 20
DEPTH_60 = 60

# K线时间区间
INTERVAL_1M = '1min'
INTERVAL_3M = '3min'
INTERVAL_5M = '5min'
INTERVAL_15M = '15min'
INTERVAL_30M = '30min'
INTERVAL_1H = '1hour'
INTERVAL_2H = '2hour'
INTERVAL_4H = '4hour'
INTERVAL_6H = '6hour'
INTERVAL_1D = 'day'
INTERVAL_3D = '3day'
INTERVAL_1W = 'week'

# 交易代码，需要后缀货币名才能完整
TRADING_SYMBOL_BTC = 'btc_'
TRADING_SYMBOL_LTC = 'ltc_'

# 委托类型
TYPE_BUY = 'buy'
TYPE_SELL = 'sell'
TYPE_BUY_MARKET = 'buy_market'
TYPE_SELL_MARKET = 'sell_market'

# 期货合约到期类型
FUTURE_EXPIRY_THIS_WEEK = 'this_week'
FUTURE_EXPIRY_NEXT_WEEK = 'next_week'
FUTURE_EXPIRY_QUARTER = 'quarter'

# 期货委托类型
FUTURE_TYPE_LONG = 1
FUTURE_TYPE_SHORT = 2
FUTURE_TYPE_SELL = 3
FUTURE_TYPE_COVER = 4

# 期货是否用现价
FUTURE_ORDER_MARKET = 1
FUTURE_ORDER_LIMIT = 0

# 期货杠杆
FUTURE_LEVERAGE_10 = 10
FUTURE_LEVERAGE_20 = 20

# 委托状态
ORDER_STATUS_NOTTRADED = 0
ORDER_STATUS_PARTTRADED = 1
ORDER_STATUS_ALLTRADED = 2
ORDER_STATUS_CANCELLED = -1
ORDER_STATUS_CANCELLING = 4


########################################################################
class OkCoinApi(object):
	"""基于Websocket的API对象"""

	#----------------------------------------------------------------------
	def __init__(self):
		"""Constructor"""
		self.apiKey = ''		# 用户名
		self.secretKey = ''	 # 密码
		self.host = ''		  # 服务器地址
		
		self.currency = ''	  # 货币类型（usd或者cny）
		
		self.ws = None		  # websocket应用对象
		self.thread = None	  # 工作线程
	
	#######################
	## 通用函数
	#######################
	
	#----------------------------------------------------------------------
	def readData(self, evt):
		"""解压缩推送收到的数据"""
		# 创建解压器
		decompress = zlib.decompressobj(-zlib.MAX_WBITS)
		
		# 将原始数据解压成字符串
		inflated = decompress.decompress(evt) + decompress.flush()
		
		# 通过json解析字符串
		data = json.loads(inflated)
		
		return data
	
	#----------------------------------------------------------------------
	def generateSign(self, params):
		"""生成签名"""
		l = []
		for key in sorted(params.keys()):
			l.append('%s=%s' %(key, params[key]))
		l.append('secret_key=%s' %self.secretKey)
		sign = '&'.join(l)
		return hashlib.md5(sign.encode('utf-8')).hexdigest().upper()
	
	#----------------------------------------------------------------------
	def onMessage(self, ws, evt):
		"""信息推送"""
		print 'onMessage'
		data = self.readData(evt)
		print data
		
	#----------------------------------------------------------------------
	def onError(self, ws, evt):
		"""错误推送"""
		print 'onError'
		print evt
		
	#----------------------------------------------------------------------
	def onClose(self, ws):
		"""接口断开"""
		print 'onClose'
		
	#----------------------------------------------------------------------
	def onOpen(self, ws):
		"""接口打开"""
		print 'onOpen'
		
	#----------------------------------------------------------------------
	def connect(self, host, apiKey, secretKey, trace=False):
		"""连接服务器"""
		self.host = host
		self.apiKey = apiKey
		self.secretKey = secretKey
		
		if self.host == OKCOIN_CNY:
			self.currency = CURRENCY_CNY
		else:
			self.currency = CURRENCY_USD
			
		websocket.enableTrace(trace)
		
		self.ws = websocket.WebSocketApp(host,
										 on_message=self.onMessage,
										 on_error=self.onError,
										 on_close=self.onClose,
										 on_open=self.onOpen)
		
		self.thread = Thread(target=self.ws.run_forever)
		self.thread.start()
		
	#----------------------------------------------------------------------
	def sendMarketDataRequest(self, channel):
		"""发送行情请求"""
		# 生成请求
		d = {}
		d['event'] = 'addChannel'
		d['binary'] = True
		d['channel'] = channel
		
		# 使用json打包并发送
		j = json.dumps(d)
		self.ws.send(j)
		
	#----------------------------------------------------------------------
	def sendTradingRequest(self, channel, params):
		"""发送交易请求"""
		# 在参数字典中加上api_key和签名字段
		params['api_key'] = self.apiKey
		params['sign'] = self.generateSign(params)
		
		# 生成请求
		d = {}
		d['event'] = 'addChannel'
		d['binary'] = True
		d['channel'] = channel
		d['parameters'] = params
		
		# 使用json打包并发送
		j = json.dumps(d)
		self.ws.send(j)
	
	#######################
	## 现货相关
	#######################
	
	#----------------------------------------------------------------------
	def subscribeSpotTicker(self, symbol):
		"""订阅现货普通报价"""
		self.sendMarketDataRequest('ok_sub_spot%s_%s_ticker' %(self.currency, symbol))
	
	#----------------------------------------------------------------------
	def subscribeSpotDepth(self, symbol, depth):
		"""订阅现货深度报价"""
		self.sendMarketDataRequest('ok_sub_spot%s_%s_depth_%s' %(self.currency, symbol, depth))
		
	#----------------------------------------------------------------------
	def subscribeSpotTradeData(self, symbol):
		"""订阅现货成交记录"""
		self.sendMarketDataRequest('ok_sub_spot%s_%s_trades' %(self.currency, symbol))
		
	#----------------------------------------------------------------------
	def subscribeSpotKline(self, symbol, interval):
		"""订阅现货K线"""
		self.sendMarketDataRequest('ok_sub_spot%s_%s_kline_%s' %(self.currency, symbol, interval))
		
	#----------------------------------------------------------------------
	def spotTrade(self, symbol, type_, price, amount):
		"""现货委托"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['type'] = str(type_)
		params['price'] = str(price)
		params['amount'] = str(amount)
		
		channel = 'ok_spot%s_trade' %(self.currency)
		
		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def spotCancelOrder(self, symbol, orderid):
		"""现货撤单"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['order_id'] = str(orderid)
		
		channel = 'ok_spot%s_cancel_order' %(self.currency)

		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def spotUserInfo(self):
		"""查询现货账户"""
		channel = 'ok_spot%s_userinfo' %(self.currency)
		
		self.sendTradingRequest(channel, {})
		
	#----------------------------------------------------------------------
	def spotOrderInfo(self, symbol, orderid):
		"""查询现货委托信息"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['order_id'] = str(orderid)
		
		channel = 'ok_spot%s_orderinfo' %(self.currency)
		
		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def subscribeSpotTrades(self):
		"""订阅现货成交信息"""
		channel = 'ok_sub_spot%s_trades' %(self.currency)
		
		self.sendTradingRequest(channel, {})
		
	#----------------------------------------------------------------------
	def subscribeSpotUserInfo(self):
		"""订阅现货账户信息"""
		channel = 'ok_sub_spot%s_userinfo' %(self.currency)
		
		self.sendTradingRequest(channel, {})
	
	#######################
	## 期货相关
	#######################
	
	#----------------------------------------------------------------------
	def subscribeFutureTicker(self, symbol, expiry):
		"""订阅期货普通报价"""
		self.sendMarketDataRequest('ok_sub_future%s_%s_ticker_%s' %(self.currency, symbol, expiry))
	
	#----------------------------------------------------------------------
	def subscribeFutureDepth(self, symbol, expiry, depth):
		"""订阅期货深度报价"""
		self.sendMarketDataRequest('ok_sub_future%s_%s_depth_%s_%s' %(self.currency, symbol,
																	  expiry, depth))
		
	#----------------------------------------------------------------------
	def subscribeFutureTradeData(self, symbol, expiry):
		"""订阅期货成交记录"""
		self.sendMarketDataRequest('ok_sub_future%s_%s_trade_%s' %(self.currency, symbol, expiry))
		
	#----------------------------------------------------------------------
	def subscribeFutureKline(self, symbol, expiry, interval):
		"""订阅期货K线"""
		self.sendMarketDataRequest('ok_sub_future%s_%s_kline_%s_%s' %(self.currency, symbol,
																	  expiry, interval))
		
	#----------------------------------------------------------------------
	def subscribeFutureIndex(self, symbol):
		"""订阅期货指数"""
		self.sendMarketDataRequest('ok_sub_future%s_%s_index' %(self.currency, symbol))
	
	#----------------------------------------------------------------------
	def futureTrade(self, symbol, expiry, type_, price, amount, order, leverage):
		"""期货委托"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['type'] = str(type_)
		params['price'] = str(price)
		params['amount'] = str(amount)
		params['contract_type'] = str(expiry)
		params['match_price'] = str(order)
		params['lever_rate'] = str(leverage)
		
		channel = 'ok_future%s_trade' %(self.currency)
		
		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def futureCancelOrder(self, symbol, expiry, orderid):
		"""期货撤单"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['order_id'] = str(orderid)
		params['contract_type'] = str(expiry)
		
		channel = 'ok_future%s_cancel_order' %(self.currency)

		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def futureUserInfo(self):
		"""查询期货账户"""
		channel = 'ok_future%s_userinfo' %(self.currency)
		
		self.sendTradingRequest(channel, {})
		
	#----------------------------------------------------------------------
	def futureOrderInfo(self, symbol, expiry, orderid, status, page, length):
		"""查询期货委托信息"""
		params = {}
		params['symbol'] = str(symbol+self.currency)
		params['order_id'] = str(orderid)
		params['contract_type'] = expiry
		params['status'] = status
		params['current_page'] = page
		params['page_length'] = length
		
		channel = 'ok_future%s_orderinfo'
		
		self.sendTradingRequest(channel, params)
		
	#----------------------------------------------------------------------
	def subscribeFutureTrades(self):
		"""订阅期货成交信息"""
		channel = 'ok_sub_future%s_trades' %(self.currency)
		
		self.sendTradingRequest(channel, {})
		
	#----------------------------------------------------------------------
	def subscribeFutureUserInfo(self):
		"""订阅期货账户信息"""
		channel = 'ok_sub_future%s_userinfo' %(self.currency)
		
		self.sendTradingRequest(channel, {})
		
	#----------------------------------------------------------------------
	def subscribeFuturePositions(self):
		"""订阅期货持仓信息"""
		channel = 'ok_sub_future%s_positions' %(self.currency)
		
		self.sendTradingRequest(channel, {})
	

